import tensorflow as tf
import os
from PIL import Image
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

os.environ["CUDA_VISIBLE_DEVICES"] = '1'


# The following functions can be used to convert a value to a type compatible
# with tf.Example.

def _bytes_feature(value):
    """Returns a bytes_list from a string / byte."""
    if isinstance(value, type(tf.constant(0))):
        value = value.numpy()  # BytesList won't unpack a string from an EagerTensor.
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _float_feature(value):
    """Returns a float_list from a float / double."""
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))


def _int64_feature(value):
    """Returns an int64_list from a bool / enum / int / uint."""
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def create_example(image_string, mask):
    feature = {
        'mask': _bytes_feature(mask),
        'image': _bytes_feature(image_string),
    }

    return tf.train.Example(features=tf.train.Features(feature=feature))


def generate(IMG_SIZE,imagelist, masks, record_file):
    """使用PIllow 压缩，转numpy的bytes"""

    writer = tf.io.TFRecordWriter(record_file)
    num = 0
    for img, mask in zip(imagelist, masks):
        # try:
        image_string = Image.open(img)
        image_string = image_string.convert('RGB')
        image_string = image_string.resize((IMG_SIZE, IMG_SIZE), Image.ANTIALIAS)
        image_string = np.asarray(image_string, np.uint8)
        image_string = image_string.tobytes()

        """masks"""
        mask = Image.open(mask).convert('L')
        mask = mask.resize((IMG_SIZE, IMG_SIZE), Image.ANTIALIAS)
        mask = np.asarray(mask, np.uint8)
        mask = mask.copy()
        """将多分类问题，转为2分类，抠图"""
        mask[mask > 0] = 255
        mask = mask.tobytes()
        tf_example = create_example(image_string, mask)
        writer.write(tf_example.SerializeToString())
        num += 1
        if num % 1000 == 0:
            print("has create {} example".format(num))
        # except Exception as e:
        #     print(e)
        #     pass


def read_tf_record(serialized_example):
    image_feature_description = {
        'mask': tf.io.FixedLenFeature([], tf.string),
        'image': tf.io.FixedLenFeature([], tf.string),
    }

    feature = tf.io.parse_example(serialized_example,
                                  features=image_feature_description)
    image = tf.io.decode_raw(feature['image'], tf.uint8)
    image = tf.reshape(image, [512, 512, 3])
    # image = tf.reverse(image, axis=[-1])  # RGB->BGR

    mask = tf.io.decode_raw(feature['mask'], tf.uint8)
    mask = tf.reshape(mask, [512, 512, 1])

    return image, mask


import glob

if __name__ == '__main__':
    root = '/hdd9/ppp/ImageSegments/ATR/humanparsing'
    images_root = root + '/JPEGImages'
    mask_root = root + '/SegmentationClassAug'
    imagels = glob.glob(images_root + '/*')
    maskls = glob.glob(mask_root + '/*')
    imagels = sorted(imagels)
    maskls = sorted(maskls)
    print(imagels[0])
    print(maskls[0])
    print(len(imagels))
    IMG_SIZE = 512
    X_train, X_test, y_train, y_test = train_test_split(imagels, maskls, test_size=0.2, random_state=42)

    generate(IMG_SIZE,X_train, y_train, root + '/record/train.record')

    generate(IMG_SIZE,X_test, y_test, root + '/record/test.record')
